import os
import requests
import time


def download_image(image_url, file_path):
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36 Edg/121.0.0.0'
    }

    attempts = 5  # 最大重试次数
    for attempt in range(attempts):
        try:
            response = requests.get(image_url, headers=headers)
            response.raise_for_status()  # 检查请求是否成功

            with open(file_path, 'wb') as file:
                file.write(response.content)

            print(f"Downloaded: {file_path}")
            return  # 成功下载，退出函数

        except Exception as e:
            print(f"Failed to download {image_url}: {e}")
            time.sleep(5)  # 等待 5 秒后重试

    print(f"Max attempts reached. Skipping {image_url}.")


def download_images(first_image_url, count):
    # 从完整的图片 URL 中提取文件名和基础 URL
    base_url = first_image_url.rsplit('/', 1)[0] + '/'  # 获取基础 URL
    initial_id = first_image_url.split('/')[-1].split('.')[0]  # 获取文件名（不带扩展名）

    # 提取基础编号（去掉最后一位字符，保留后面的编号）
    base_id = initial_id[:-1]  # 例如，15673714223053

    # 创建以基础编号命名的文件夹
    folder_path = base_id  # 使用基础编号作为文件夹名称
    os.makedirs(folder_path, exist_ok=True)  # 创建目录，如果已存在则不抛出错误

    for i in range(1, count + 1):  # 从 1 开始编号
        # 根据规则生成图片 ID
        image_id = f"{base_id}{i}"  # 将基础 ID 加上序号

        image_url = f"{base_url}{image_id}.jpg"

        # 文件路径
        file_path = os.path.join(folder_path, f"{image_id}.jpg")

        # 如果文件已经存在，直接跳过
        if os.path.exists(file_path):
            print(f"Image {file_path} already exists. Skipping.")
            continue

            # 下载图片
        download_image(image_url, file_path)


if __name__ == "__main__":
    # 输入完整的第一张图片的 URL 和下载数量
    first_image_url = input("请输入第一张图片的完整 URL: ")
    num_images = int(input("请输入要下载的图片数量: "))

    download_images(first_image_url, num_images)